#include "signal_server.hpp"  // 引入信令服务器头文件
#include <iostream>  // 引入输入输出流
#include <asio.hpp>  // 引入 ASIO 库
#include <unordered_map>  // 引入 unordered_map，用于存储注册信息
#include <nlohmann/json.hpp>  // 引入 JSON 库

using asio::ip::tcp;  // 使用 TCP 协议
using json = nlohmann::json;  // 定义 JSON 类型

// 构造函数：初始化 I/O 上下文和端口号
SignalServer::SignalServer(asio::io_context& io_context, short port)
    : io_context_(io_context), acceptor_(io_context, tcp::endpoint(tcp::v4(), port)) {}

// 运行信令服务器：监听端口并接受连接
void SignalServer::run() {
    std::cout << "Signal server is running on port 9000..." << std::endl;
    while (true) {
        tcp::socket socket(io_context_);
        acceptor_.accept(socket);  // 接受客户端连接
        // 启动新线程处理每个客户端
        std::thread(&SignalServer::handle_client, this, std::move(socket)).detach();
    }
}

// 处理客户端请求：读取请求并决定注册或查询操作
void SignalServer::handle_client(tcp::socket socket) {
    try {
        asio::streambuf buf;  // 用于存储客户端请求
        asio::read_until(socket, buf, "\n");  // 读取直到换行符为止
        std::istream input(&buf);  // 创建输入流
        std::string request;
        std::getline(input, request);  // 获取客户端请求

        json j = json::parse(request);  // 解析 JSON 请求
        std::string type = j["type"];  // 获取请求类型

        // 根据请求类型调用不同的处理函数
        if (type == "register") {
            register_client(j, socket);  // 处理注册请求
        } else if (type == "query") {
            query_client(j, socket);  // 处理查询请求
        }
    } catch (...) {
        std::cerr << "Client disconnected or error occurred." << std::endl;
    }
}

// 处理客户端注册：将 UUID 和 IP:Port 存入注册表
void SignalServer::register_client(const json& request, tcp::socket& socket) {
    std::string uuid = request["uuid"];  // 获取客户端的 UUID
    auto remote_endpoint = socket.remote_endpoint();  // 获取客户端的远程 IP:Port
    uuid_registry_[uuid] = remote_endpoint;  // 存储在注册表中

    std::cout << "Registered: " << uuid << " at " << remote_endpoint << std::endl;

    // 发送注册成功的响应
    json response = {{"status", "registered"}};
    std::string response_str = response.dump() + "\n";
    asio::write(socket, asio::buffer(response_str));  // 发送响应
}

// 处理客户端查询：返回目标 UUID 的 IP:Port 信息
void SignalServer::query_client(const json& request, tcp::socket& socket) {
    std::string target_uuid = request["target"];  // 获取目标 UUID
    json response;

    // 查询注册表，找到目标 UUID 的 IP:Port
    if (uuid_registry_.find(target_uuid) != uuid_registry_.end()) {
        auto ep = uuid_registry_[target_uuid];  // 获取目标 IP:Port
        response["ip"] = ep.address().to_string();  // 返回 IP 地址
        response["port"] = ep.port();  // 返回端口号
    } else {
        response["error"] = "UUID not found";  // 未找到 UUID
    }

    // 发送查询结果
    std::string response_str = response.dump() + "\n";
    asio::write(socket, asio::buffer(response_str));

    // 新增：如果查询是打洞请求，返回给客户端打洞信息
    if (request.contains("punchthrough") && request["punchthrough"] == true) {
        json punchthrough_response = {{"status", "punchthrough"}};
        std::string punchthrough_response_str = punchthrough_response.dump() + "\n";
        asio::write(socket, asio::buffer(punchthrough_response_str));
    }
}
